import bisect
import math
import sys
import warnings

import pylab
import matplotlib
from matplotlib import pyplot
from matplotlib.colors import LogNorm
import numpy

from histogram.h1d import Histogram1D
from histogram.mplif import MPLInterface
from histogram import decorators

class Histogram2D(MPLInterface):
    _rcparams = {'axes.grid': True,
                 'axes.labelsize': 'x-large',
                 'axes.titlesize': 'x-large',
                 'axes.formatter.limits': (-3, 3),
                 'font.size': 24.0,
                 'figure.figsize': (8, 6),
                 'grid.color': 'w',
                 'text.latex.preamble': [r'\usepackage{amsmath}'],
                 'text.usetex': True,
                 'xtick.labelsize': 'large',
                 'xtick.direction': 'out',
                 'ytick.labelsize': 'large',
                 'ytick.direction': 'out',
                 }


    """
    Two-dimensional histogram object.
    
    Parameters
    ----------
    name   : string
        A name for the histogram
    nbins  : sequence of length 2
        (``number of bins in x``, ``number of bins in y``)
    arange : a sequence of two sequences of length 2
        ((``x axis minimum``, ``x axis maximum``), 
         (``y axis minimum``, ``y axis maximum``))
    xlabel : string
        Label for the x-axis
    ylabel : string
        Label for the y-axis
    xunits : string 
        Units for the x-axis
    yunits : string 
        Units for the y-axis

    Optional Parameters
    -------------------
    title : string 
        Title of the histogram (default is "`xlabel` vs `ylabel`")
    fontsize  :  float | {'xx-small' , 'x-small' , 'small' , 'medium' , 
                          'large' , 'x-large' , 'xx-large'}
        Size of text when drawn
    ext       : {'png' , 'gif' , 'jpg'}
        Default file extension to add to name when histogram is drawn 
        and saved to file
    bufsize   : int
        length of internal data buffer; larger bufsizes are faster 
        for large datasets but take up more memory
    assoc_cols     : 2-tuple of string
        Two column names (x and y) which are associated with this histogram
    selection : string
        A PyTables-style selection which is associated with this histogram
    logz      : bool
        Set the histogram to have a logarithmic z-axis when drawn            
    use_latex : True | False
        Turn on/off Matplotlib feeding image text through LaTeX

    """

    def __init__(self, name, nbins, aranges, xlabel, ylabel, xunits, yunits,
                 title='', fontsize='x-large', 
                 entries=dict(x=1.25, y=-0.2),
                 bufsize=1000000, assoc_cols=[], logz=False, 
                 bins=None, usetex=True):

        self.name      = name
        self.nxbins    = nbins[0]
        self.nybins    = nbins[1]
        self.xrange    = aranges[0]
        self.yrange    = aranges[1]
        self.assoc_cols = assoc_cols

        self._i        = 0;
        self._tot_i    = 0;

        self._bufsize  = bufsize

        self.low_edges_x = None
        self.low_edges_y = None
        self.bin_width = (float(self.xrange[1] - self.xrange[0]) / self.nxbins,
                          float(self.yrange[1] - self.yrange[0]) / self.nybins)


        if bins is not None:
            self._bins  = bins
            self.low_edges_x, bin_width_x = numpy.linspace(self.xrange[0], 
                                                           self.xrange[1], 
                                                           self.nxbins+1, 
                                                           retstep=True)
            self.low_edges_y, bin_width_y = numpy.linspace(self.yrange[0], 
                                                           self.yrange[1], 
                                                           self.nybins+1, 
                                                           retstep=True)

            self._mean  = self._mean_from_bins()
            self._var   = self._var_from_bins()
            self._xmin, self._xmax = self._min_from_bins()
            self._ymin, self._ymax = self._max_from_bins()
            self._tot_i = numpy.sum(bins)

        else:
            self._bins  = None 
            self._mean   = 0.0 
            self._var    = 0.0
            self._xmin    = 1e250
            self._xmax    = -1e250
            self._ymin    = 1e250
            self._ymax    = -1e250

        self._buf  = None
        self._wbuf = None 

        if title is None:
            title = '{} vs {}'.format(ylabel, xlabel)

        super(Histogram2D, self).__init__(rcparams=self._rcparams,
                                          axesprops=(0.185, 0.165, 0.8, 0.72),
                                          xlabel=self._make_axlabel(xlabel,
                                                                    xunits),
                                          ylabel=self._make_axlabel(ylabel,
                                                                    yunits),
                                          xunits=xunits,
                                          yunits=yunits,
                                          title=title,
                                          logz=logz,
                                          entries=entries,
                                          usetex=usetex,
                                          colorbar=True)


    def _setup_arrays(self):
        if self._bins is None:
            self._bins  = numpy.zeros((self.nxbins, self.nybins))

        if self._buf is None:
            self._buf   = numpy.zeros((self._bufsize, 2))

        if self._wbuf is None:
            self._wbuf  = numpy.ones((self._bufsize,))

        if self.low_edges_x is None:
            self.low_edges_x, bin_width_x = numpy.linspace(self.xrange[0], 
                                                           self.xrange[1], 
                                                           self.nxbins+1, 
                                                           retstep=True)

        if self.low_edges_y is None:
            self.low_edges_y, bin_width_y = numpy.linspace(self.yrange[0], 
                                                           self.yrange[1], 
                                                           self.nybins+1, 
                                                           retstep=True)
        
    
    @decorators.draw_validate
    def draw(self, **kwargs):
        """
        Prepare a canvas for showing via GUI or saving to an image file.

        Optional Parameters
        -------------------
        cmap      : {'autumn', 'bone', 'cool', 'copper', 'flag', 'gray', 
                     'hot', 'hsv', 'jet', 'pink', 'prism', 'spring', 
                     'summer', 'winter', 'spectral'}
            Use colormap ``cmap``
        ext       : {'png' , 'gif' , 'jpg'}
            File extension to add to name when histogram is drawn and saved 
            to file
        logz      : bool
            Set the histogram to have a logarithmic z-axis when drawn
        name   : string
            A name for the histogram
        xlabel : string
            Label for the x-axis
        xunits : string 
            Units for the x-axis
        ylabel : string
            Label for the y-axis
        yunits : string 
            Units for the y-axis
        """
        
        super(Histogram2D, self).draw(**kwargs)

        kwargs = self._mplif_curkwargs

        if kwargs['entries']:
            pyplot.text(kwargs['entries']['x'], kwargs['entries']['y'],
                        'Entries: {0:g}'.format(self.nentries()),
                        transform=kwargs['axes'].transAxes,
                        fontsize=matplotlib.rcParams['font.size'],
                        horizontalalignment='right',
                        verticalalignment='bottom')
                

        pyplot.axis(xmin=self.xrange[0], xmax=self.xrange[1], 
                    ymin=self.yrange[0], ymax=self.yrange[1])
        
        if 'cmap' not in kwargs:
            kwargs['cmap'] = 'jet'

        try:
            pyplot.set_cmap(getattr(pylab.cm, kwargs['cmap']))
        except (AttributeError, TypeError):
            msg = "Histogram2D.draw() failed: no colormap `{}' in pyplot"
            warnings.warn(msg.format(kwargs['cmap']))
            return kwargs['axes']

        return kwargs['axes']
    

    def _mplif_plot(self, mplif_kw, other_kw):
        X, Y = numpy.meshgrid(self.low_edges_x, self.low_edges_y)

        if 'logz' in mplif_kw and mplif_kw['logz']:
            bins = self._bins.T + 1
            pyplot.pcolormesh(X, Y, bins, 
                              norm=LogNorm(vmin=bins.min(), vmax=bins.max()))
        else:
            pyplot.pcolormesh(X, Y, self._bins.T)
        
        if 'colorbar' in mplif_kw and mplif_kw['colorbar']:
            pyplot.colorbar()


    def fill(self, value, weight=1, loud=False):
        """
        Fill a histogram with values.
                
        Parameters
        ----------
        
        value : ``(x, y)`` | sequence of ``(x, y)``
            Values to be binned

        Optional Parameters
        -------------------
        weight : ``(x_weight, y_weight)`` | sequence of ``(x_weight, y_weight)``
            Weights to apply when binning values
        loud   : bool
            Print a progress indicator when binning multiple values
        """

        if self._buf is None:
            self._setup_arrays()

        try:
            self._buf[self._i] = value
        except IndexError:
            self.flush()
            try:
                self._buf[self._i] = value
            except ValueError:
                return self._fill_array(numpy.asarray(value), weight, loud)
        except ValueError:
            return self._fill_array(numpy.asarray(value), weight, loud)

        self._wbuf[self._i] = weight
        self._i += 1

    def _fill_array(self, values, weights, loud):                       
        if not len(values):
            return

        valshape = numpy.shape(values)

        if len(valshape) != 2 or 2 not in valshape:
            msg = "Trying to fill 2D histogram with array of shape `{0}'; " + \
                "valid shapes are (2,n) or (n,2)"
            raise ValueError(msg.format(valshape))
        
        if valshape[1] != 2:
            try:
                values = values.T
            except AttributeError:
                values = numpy.asarray(values).T

        try:
            x = iter(weights)
        except TypeError:
            pass
        else:
            return self._fill_array_weights(values, weights)
        
        self.flush()
        
        val_len = len(values)

        if val_len < self._bufsize:
            self._i = val_len
            self._buf[:self._i] = values[:self._i]
            return

        for i in xrange(0, val_len, self._bufsize):
            if loud:
                msg = "Histogram1D `{0}': {1}/{2} events binned\r"
                sys.stderr.write(msg.format(self.name, i, val_len))
                sys.stderr.flush()
            val_arr = values[i : i + self._bufsize]
            self._i = len(val_arr)
            self._buf[:self._i]  = val_arr
            self.flush()

    def _fill_array_weights(self, values, weights):
        self.flush()
        
        val_len  = len(values)
        if val_len < self._bufsize:
            self._i = len(values)
            self._buf[:self._i]  = values
            self._wbuf[:self._i] = weights
            return

        for i in xrange(0, val_len, self._bufsize):
            val_arr = values[i : i + self._bufsize]
            self._i = len(val_arr)
            self._buf[:self._i]  = val_arr
            self._wbuf[:self._i] = weights[i : i + self._bufsize]
            self.flush()
    

    def clear(self):
        """Set all bin contents to zero."""

        if self._buf is not None:
            self._buf   = numpy.zeros((self._bufsize,))

        if self._bins is not None:
            self._bins  = numpy.zeros((self.nxbins,self.nybins))

        self._mean   = 0.0 
        self._var    = 0.0
        self._xmin    = 1e250
        self._xmax    = -1e250
        self._ymin    = 1e250
        self._ymax    = -1e250
        self._i      = 0
        self._tot_i  = 0


    def flush(self):
        """
        Bin events in buffer and clear the buffer.
        """

        if self._buf is None or self._bins is None:
            self._setup_arrays()

        if self._i == 0:
            return
        
        buf_to_flush = self._buf[:self._i]

        bufxmax = max(buf_to_flush.T[0])
        if bufxmax > self._xmax:
            self._xmax = bufxmax
        bufxmin = min(buf_to_flush.T[0])
        if bufxmin < self._xmin:
            self._xmin = bufxmin

        bufymax = max(buf_to_flush.T[1])
        if bufymax > self._ymax:
            self._ymax = bufymax
        bufymin = min(buf_to_flush.T[1])
        if bufymin < self._ymin:
            self._ymin = bufymin

        
        self._mean = (self._mean * self._tot_i + numpy.mean(buf_to_flush) * 
                      self._i) / (self._tot_i + self._i)
        self._var  += sum((buf_to_flush - self._mean)**2)
 
        bufbins, edges = numpy.histogramdd(buf_to_flush,
                                           bins=(self.low_edges_x, 
                                                 self.low_edges_y),
                                           weights=self._wbuf[:self._i],
                                           range=(self.xrange, self.yrange),
                                           )
        
        self._bins += bufbins
        self._tot_i += self._i
        self._i = 0
        
    def project_to_y(self, x_start=0, x_stop=None):
        """
        Project the histogram onto the x-axis

        Optional Parameters
        -------------------
        bin_start : int
            Bin index of lower bound of projection
        bin_stop  : int
            Bin index of upper bound of projection

        Returns
        -------
        a Histogram1D object.
        """

        if self._bins is None:
            self._setup_arrays()
        else:
            self.flush()
        
        bin_start, iy = self.bin_index(x_start, 0)
  
        if x_stop is None:
            bin_stop = self._bins.shape[0]
        else:
            bin_stop, iy = self.bin_index(x_stop, 0)

        bins = numpy.sum(self._bins[bin_start:bin_stop], 0)
        return Histogram1D(self.name + "_py", len(bins),
                           self.yrange,
                           self.ylabel,
                           self.yunits,
                           bins=bins)

    def project_to_x(self, bin_start=0, bin_stop=None):
        """
        Project the histogram onto the x-axis

        Optional Parameters
        -------------------
        bin_start : int
            Bin index of lower bound of projection
        bin_stop  : int
            Bin index of upper bound of projection

        Returns
        -------
        a Histogram1D object.

        """

        if self._bins is None:
            self._setup_arrays()
        else:
            self.flush()

        if bin_stop is None:
            bin_stop = self._bins.shape[1]

        bins = numpy.sum(self._bins.T[bin_start:bin_stop], 0)
        return Histogram1D(self.name + "_px", len(bins),
                           self.xrange,
                           self.xlabel,
                           self.xunits,
                           bins=bins)


    
    def set_buffer_size(self, bufsize):
        """
        Set internal buffer size to new value.
        
        Parameters
        ----------
        
        bufsize : int
            desired new buffer size
        """

        self._bufsize = bufsize

        if self._buf is not None:            
            self.flush()
            
            self._buf     = numpy.zeros((self._bufsize,))
            self._wbuf    = numpy.ones((self._bufsize,))


    def __getstate__(self):
        self.flush()

        outdict = {}
        outdict.update(self.__dict__)
        
        outdict['low_edges_x'] = None
        outdict['low_edges_y'] = None
        outdict['_bins'] = None
        outdict['_buf'] = None
        outdict['_wbuf'] = None
        
        return outdict

    def _min_from_bins(self):
        return (0,0)

    def _max_from_bins(self):
        return (0,0)

    def _mean_from_bins(self):
    #   1D case
    #   mean = sum([self._bins[i]*(self.low_edges[i] + (self.bin_width/2.0)) for i in xrange(self.nbins)]) / sum(self._bins)
        return 0.0

    def _var_from_bins(self):
    #   1D case
    #   var = sum([self._bins[i]*(self.low_edges[i] + (self.bin_width/2.0) - self._mean)**2 for i in xrange(self.nbins)])
        return 0.0


    def bins(self):
        """Returns a reference to the array of bin contents."""

        if self._bins is None:
            self._setup_arrays()
        else:
            self.flush()

        return self._bins

    def max_bin_index(self):
        """
        Returns the indices ``(ix, iy)`` of the bin with the most counts.
        """

        assert self._bins is not None, "Histogram is empty"

        self.flush()

        flat_index = numpy.argmax(self._bins)

        return numpy.unravel_index(numpy.argmax(self._bins), self._bins.shape)

    def max_bin_content(self):
        """
        Returns the number of counts in the bin with the most counts.
        """

        assert self._bins is not None, "Histogram is empty"

        i, j = self.max_bin_index()
        return self._bins[i][j]

    def max_bin_xy(self):
        """
        Returns the x-y coordinates ``(x, y)`` of the bin with the most counts.
        """

        assert self._bins is not None, "Histogram is empty"

        return self.bin_center(*self.max_bin_index())

    def bin_center(self, ix, iy):
        """
        Returns the x-y coordinates ``(x, y)`` of the center 
        of the bin at ``(ix, iy)``.
        """

        if self.low_edges_x is None or self.low_edges_y is None:
            self._setup_arrays()

            return (self.low_edges_x[ix] + self.bin_width[0]*0.5, 
                    self.low_edges_y[iy] + self.bin_width[1]*0.5)

    def bin_index(self, x, y):
        """
        Return the indices ``(ix, iy)`` of the bin that ``(x, y)`` 
        would be placed into.
        """

        if self.low_edges_x is None or self.low_edges_y is None:
            self._setup_arrays()
        
        return (bisect.bisect(self.low_edges_x, x), 
                bisect.bisect(self.low_edges_y, y))
    
    def data_xmax(self):
        """
        Returns the current maximum x-value binned.
        """

        assert self._bins is not None, "Histogram is empty"

        self.flush()
        return self._xmax

    def data_xmin(self):
        """
        Returns the current minimum x-value binned.
        """

        assert self._bins is not None, "Histogram is empty"

        self.flush()
        return self._xmin

    def data_ymax(self):
        """
        Returns the current maximum y-value binned.
        """

        assert self._bins is not None, "Histogram is empty"

        self.flush()
        return self._ymax

    def data_ymin(self):
        """
        Returns the current minimum y-value binned.
        """

        assert self._bins is not None, "Histogram is empty"

        self.flush()
        return self._ymin

    def data_mean(self):
        """
        Returns the mean ``(mu_x, mu_y)`` of all of the values binned.
        """

        assert self._bins is not None, "Histogram is empty"

        self.flush()
        return self._mean

    def data_var(self):
        """
        Returns the variance ``(var_x, var_y)`` of all of the values binned.
        """

        assert self._bins is not None, "Histogram is empty"

        self.flush()
        return self._var

    def data_std(self):
        """
        Returns the standard deviation ``(std_x, std_y)`` of 
        all of the values binned.
        """

        assert self._bins is not None, "Histogram is empty"

        self.flush()
        return math.sqrt(self._var / float(self._tot_i))

    def nentries(self):
        """
        Returns the total number of values binned.
        """

        if self._bins is None:
            return 0
        
        self.flush()
        return self._tot_i

        counts = numpy.sum(self._bins)

        if counts < int(counts) > 0.0001:
            return counts
        else:
            return int(counts)


    def __len__(self):
        return self.nentries()

                    
    def __iadd__(self, rhs):     
        if self._bins is None:
            self._setup_arrays()

        if rhs._bins is None:
            rhs._setup_arrays()

        self.flush()
        rhs.flush()

        if len(rhs) == 0:
            return self
        
        tot_entries = self._tot_i + rhs._tot_i
        self._mean = self._mean * (self._tot_i / float(tot_entries)) \
            + rhs._mean * (rhs._tot_i / float(tot_entries))
        self._var  = self._var * (self._tot_i / float(tot_entries)) \
            + rhs._var * (rhs._tot_i / float(tot_entries))

        self._tot_i = tot_entries
        
        if self._xmax < rhs._xmax:
            self._xmax = rhs._xmax
        
        if self._xmin > rhs._xmin:
            self._xmin = rhs._xmin        

        if self._ymax < rhs._ymax:
            self._ymax = rhs._ymax
        
        if self._ymin > rhs._ymin:
            self._ymin = rhs._ymin        

        self._bins += rhs._bins
        
        return self

    def _make_axlabel(self, lbl, un):
        if un is not None:
            new_label = '{} ({})'.format(lbl, un)
        else:
            new_label = lbl

        return new_label
    
    def __repr__(self):
        return str(self)

    def __str__(self):
        s = "Histogram2D `{0}'; {1} bins from {2} to {3} {4} in X; " \
            + "{5} bins from {6} to {7} {8} in Y; {9} entries"
        return s.format(self.name, self.nxbins, self.xrange[0], 
                        self.yrange[1], self.xunits, 
                        self.nybins, self.yrange[0], 
                        self.yrange[1], self.yunits,
                        self.nentries())
